package com.chromatic.competition.service;

import android.os.AsyncTask;

import org.json.JSONArray;
import org.json.JSONException;
import org.json.JSONObject;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.HttpURLConnection;
import java.net.URL;
import java.nio.charset.StandardCharsets;

public class StabilityAIService {

    private static final String API_URL = "https://api.stability.ai";
    private static final String ENGINE_ID = "stable-diffusion-xl-beta-v2-2-2";
//    private static final String ENGINE_ID = "stable-diffusion-768-v2-1";
//    private static final String ENGINE_ID = "stable-diffusion-v1-5";
    private static final String API_KEY = "sk-yJXtFnYPPmeoC84QEtejAThtwI2ZnLDBeIxH5K2nUAkUdl9B";

    public interface ImageGenerationListener {
        void onImageGenerated(String base64Image);
        void onError(String message);
    }

    /**
     * 生成图像
     *
     * @param text     文本内容
     * @param listener 图像生成监听器
     */
    public static void generateImage(String text, ImageGenerationListener listener) {
        new GenerateImageTask(listener).execute(text);
    }

    private static class GenerateImageTask extends AsyncTask<String, Void, String> {

        private final ImageGenerationListener listener;

        public GenerateImageTask(ImageGenerationListener listener) {
            this.listener = listener;
        }

        @Override
        protected String doInBackground(String... params) {
            String text = params[0];
            String base64Image = "";

            try {
                URL url = new URL(API_URL + "/v1/generation/" + ENGINE_ID + "/text-to-image");
                HttpURLConnection conn = (HttpURLConnection) url.openConnection();
                conn.setRequestMethod("POST");
                conn.setRequestProperty("Content-Type", "application/json");
                conn.setRequestProperty("Accept", "application/json");
                conn.setRequestProperty("Authorization", "Bearer " + API_KEY);
                conn.setDoOutput(true);

                JSONObject requestBody = new JSONObject();
                requestBody.put("text_prompts", new JSONArray().put(new JSONObject().put("text", text)));
                requestBody.put("cfg_scale", 7);
                requestBody.put("clip_guidance_preset", "FAST_BLUE");
                requestBody.put("height", 512);
                requestBody.put("width", 512);
                requestBody.put("samples", 1);
                requestBody.put("steps", 100);

                OutputStream outputStream = conn.getOutputStream();
                outputStream.write(requestBody.toString().getBytes(StandardCharsets.UTF_8));
                outputStream.close();

                int responseCode = conn.getResponseCode();

                if (responseCode == HttpURLConnection.HTTP_OK) {
                    String response = readResponse(conn.getInputStream());
                    JSONObject responseJson = new JSONObject(response);
                    JSONArray artifacts = responseJson.getJSONArray("artifacts");

                    if (artifacts.length() > 0) {
                        JSONObject artifact = artifacts.getJSONObject(0);
                        base64Image = artifact.getString("base64");
                    } else {
                        listener.onError("No image data found in response");
                    }
                } else {
                    String errorResponse = readResponse(conn.getErrorStream());
                    listener.onError("Error: " + errorResponse);
                }

                conn.disconnect();
            } catch (IOException | JSONException e) {
                e.printStackTrace();
                listener.onError("Error: " + e.getMessage());
            }

            return base64Image;
        }

        @Override
        protected void onPostExecute(String base64Image) {
            listener.onImageGenerated(base64Image);
        }
    }

    /**
     * 读取响应内容
     *
     * @param inputStream 输入流
     * @return 响应内容字符串
     * @throws IOException IO异常
     */
    private static String readResponse(InputStream inputStream) throws IOException {
        BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream));
        StringBuilder response = new StringBuilder();
        String line;
        while ((line = reader.readLine()) != null) {
            response.append(line);
        }
        reader.close();
        return response.toString();
    }
}